import os
import operator
from dotenv import load_dotenv
from typing import Annotated, TypedDict

from langchain_core.messages import AnyMessage, HumanMessage
# 我们不再需要 DuckDuckGo，而是需要 @tool 装饰器
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI

from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode, tools_condition
from langgraph.checkpoint.memory import MemorySaver

# --- 1. 环境准备 ---
load_dotenv()

# a. 定义一个新的、可靠的本地工具
@tool
def python_calculator(expression: str) -> str:
    """
    一个可以执行 Python 数学表达式的计算器。
    例如: '3*4+5' or '2**8'
    """
    try:
        result = eval(expression, {"__builtins__": None}, {})
        return f"表达式 '{expression}' 的计算结果是: {result}"
    except Exception as e:
        return f"执行表达式 '{expression}' 时出错: {e}"

# b. 将新工具放入工具列表
tools = [python_calculator]

# --- 2. 状态定义 ---
# 与上一个例子相同，我们用 Annotated 来维护一个不断增长的消息列表。
class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], operator.add]

# --- 3. 节点定义 ---
# 初始化模型并绑定工具
llm = ChatOpenAI(model="qwen-plus-latest", base_url=os.getenv("OPENAI_BASE_URL"))
model_with_tools = llm.bind_tools(tools)

# a. Agent 节点: 代理的核心大脑，决定下一步行动
def agent_node(state: AgentState):
    print("---AGENT: 思考中...---")
    response = model_with_tools.invoke(state["messages"])
    return {"messages": [response]}

# b. Tool 节点: 执行工具
tool_node = ToolNode(tools)

# --- 4. 图的构建 ---
workflow = StateGraph(AgentState)

workflow.add_node("agent", agent_node)
workflow.add_node("tools", tool_node)

# 条件边逻辑被 langgraph 内置的 `tools_condition` 取代。
# `tools_condition` 的作用和我们之前写的 `should_continue` 完全一样：
# 检查最新的消息，如果包含 tool_calls，就路由到 'tools' 节点，否则就结束 (END)。
# 这让代码更简洁。
workflow.add_conditional_edges(
    "agent",
    tools_condition,
)

# 普通边：工具执行完后，结果返回给 agent 重新思考
workflow.add_edge("tools", "agent")

# 设置入口点
workflow.set_entry_point("agent")

# --- 5. 添加中断点并编译 ---
# `interrupt_before=["tools"]` 是关键！
# 这行代码告诉 langgraph：在每次要进入 'tools' 节点之前，暂停执行。
# 我们还需要一个 checkpointer 来保存图的状态，以便在中断后可以恢复。
checkpointer = MemorySaver()
app = workflow.compile(checkpointer=checkpointer, interrupt_before=["tools"])

# --- 6. 人机交互运行 ---
if __name__ == "__main__":
    # c. 更新问题以触发新工具
    question = "计算一下 (351 * 1.15 - 200) / 2 的结果是多少？"
    
    # 我们需要为每次独立的对话提供一个唯一的、可配置的ID（比如 session_id）。
    # checkpointer 会使用这个ID来区分不同的对话存档。
    config = {"configurable": {"thread_id": "user_session_1"}}
    inputs = {"messages": [HumanMessage(content=question)]}
    
    # `app.invoke` 会运行图，直到遇到中断点或结束。
    # 因为我们设置了中断，它会在调用工具前停下来。
    app.invoke(inputs, config)
    
    # 使用 get_state 获取当前的状态快照
    snapshot = app.get_state(config)

    # 如果 .next 属性不为空，说明图已暂停，等待进入下一个节点
    if snapshot.next:
        print("\n---代理请求使用工具---")
        # 我们可以检查代理的计划
        print("下一步将要执行的节点:", snapshot.next)
        last_message = snapshot.values["messages"][-1]
        print("工具调用请求详情:", last_message.tool_calls)
        
        user_approval = input("\n您是否批准执行该工具？ (yes/no): ")
        
        if user_approval.lower() == "yes":
            print("\n---用户批准，继续执行---")
            # `app.invoke` 传入 None，从中断的地方继续执行
            final_result = app.invoke(None, config=config)
            print("\n---最终结果---")
            final_result["messages"][-1].pretty_print()
        else:
            print("\n---用户拒绝，流程终止---")
    else:
        # 如果代理不需要使用工具，直接就结束了
        print("\n---代理未请求使用工具，直接输出结果---")
        # 我们可以从状态中获取最终消息
        final_message = app.get_state(config).values["messages"][-1]
        final_message.pretty_print() 